from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.sql.types import *
import random
from datetime import datetime, timedelta

spark = SparkSession.builder.appName("MockData").getOrCreate()


# 生成模拟用户数据
def generate_mock_users(n=1000):
    genders = ["male", "female", "unknown"]
    occupations = ["IT", "Finance", "Education", "Healthcare", "Retail", "Manufacturing", "Other"]
    educations = ["High School", "Bachelor", "Master", "PhD", "Other"]
    marital_status = ["Single", "Married", "Divorced", "Unknown"]
    income_levels = ["Low", "Medium", "High", "Very High"]
    provinces = ["Guangdong", "Jiangsu", "Zhejiang", "Shandong", "Henan", "Sichuan", "Hubei", "Other"]
    city_levels = ["Tier 1", "Tier 2", "Tier 3", "Tier 4", "Other"]

    users = []
    for i in range(1, n + 1):
        user_id = f"user_{i}"
        gender = random.choice(genders)
        age = random.randint(18, 70)
        occupation = random.choice(occupations)
        education = random.choice(educations)
        marital = random.choice(marital_status)
        has_children = random.random() > 0.5 if marital == "Married" else False
        income = random.choice(income_levels)
        province = random.choice(provinces)
        city = f"City_{random.randint(1, 30)}"
        city_level = random.choice(city_levels)

        users.append(Row(
            user_id=user_id,
            gender=gender,
            age=age,
            occupation=occupation,
            education=education,
            marital_status=marital,
            has_children=has_children,
            income_level=income,
            province=province,
            city=city,
            city_level=city_level
        ))

    return spark.createDataFrame(users)


# 生成模拟行为数据
def generate_mock_behaviors(users_df, n_per_user=10):
    behavior_types = ["click", "view", "collect", "add_cart", "purchase"]
    categories = [f"category_{i}" for i in range(1, 21)]

    behaviors = []
    user_ids = [row.user_id for row in users_df.collect()]

    for user_id in user_ids:
        n_behaviors = random.randint(1, n_per_user)
        for i in range(n_behaviors):
            item_id = f"item_{random.randint(1, 1000)}"
            category_id = random.choice(categories)
            behavior_type = random.choice(behavior_types)
            days_ago = random.randint(0, 90)
            behavior_time = datetime.now() - timedelta(days=days_ago, hours=random.randint(0, 23),
                                                       minutes=random.randint(0, 59))
            duration = random.randint(1, 300) if behavior_type in ["view", "click"] else 0
            page_depth = random.randint(1, 10) if behavior_type in ["view", "click"] else 0
            is_purchase = behavior_type == "purchase"
            order_id = f"order_{random.randint(1, 10000)}" if is_purchase else None

            behaviors.append(Row(
                user_id=user_id,
                item_id=item_id,
                category_id=category_id,
                behavior_type=behavior_type,
                behavior_time=behavior_time,
                duration=duration,
                page_depth=page_depth,
                is_purchase=is_purchase,
                order_id=order_id
            ))

    return spark.createDataFrame(behaviors)


# 生成模拟订单数据
def generate_mock_orders(users_df, behaviors_df):
    payment_methods = ["Alipay", "WeChat Pay", "Credit Card", "Debit Card", "Cash"]
    order_statuses = ["Paid", "Shipped", "Completed", "Cancelled", "Refunded"]

    # 筛选出购买行为
    purchase_behaviors = behaviors_df.filter(behaviors_df.is_purchase == True).collect()

    orders = []
    order_items = []

    for behavior in purchase_behaviors:
        order_id = behavior.order_id
        user_id = behavior.user_id
        order_time = behavior.behavior_time
        payment_time = order_time + timedelta(minutes=random.randint(1, 30))
        delivery_time = payment_time + timedelta(days=random.randint(1, 5))
        receive_time = delivery_time + timedelta(days=random.randint(1, 7))

        # 随机订单状态
        status = random.choice(order_statuses)
        if status == "Cancelled":
            receive_time = None
            delivery_time = None
        elif status == "Shipped":
            receive_time = None
        elif status == "Paid":
            receive_time = None
            delivery_time = None

        # 订单金额
        order_amount = round(random.uniform(10, 1000), 2)
        discount_rate = random.uniform(0, 0.3)
        discount_amount = round(order_amount * discount_rate, 2)
        actual_payment = order_amount - discount_amount

        payment_method = random.choice(payment_methods)

        orders.append(Row(
            order_id=order_id,
            user_id=user_id,
            order_time=order_time,
            payment_time=payment_time,
            delivery_time=delivery_time,
            receive_time=receive_time,
            order_amount=order_amount,
            discount_amount=discount_amount,
            actual_payment=actual_payment,
            payment_method=payment_method,
            order_status=status
        ))

        # 生成订单商品明细
        n_items = random.randint(1, 5)
        for i in range(n_items):
            item_id = f"item_{random.randint(1, 1000)}"
            category_id = f"category_{random.randint(1, 20)}"
            brand_id = f"brand_{random.randint(1, 50)}"
            item_name = f"Product {item_id}"
            item_price = round(random.uniform(5, 500), 2)
            quantity = random.randint(1, 10)
            item_amount = item_price * quantity

            # 随机是否退货
            is_return = random.random() < 0.1
            return_time = (receive_time + timedelta(days=random.randint(1, 30))) if is_return and receive_time else None

            order_items.append(Row(
                order_id=order_id,
                item_id=item_id,
                category_id=category_id,
                brand_id=brand_id,
                item_name=item_name,
                item_price=item_price,
                quantity=quantity,
                item_amount=item_amount,
                is_return=is_return,
                return_time=return_time
            ))

    return spark.createDataFrame(orders), spark.createDataFrame(order_items)

# 生成模拟营销活动响应数据
def generate_mock_marketing_responses(users_df):
    campaign_types = ["Promotion", "New Product", "Brand Awareness", "Retargeting"]
    conversion_types = ["Purchase", "Add to Cart", "Sign Up", "Download", "None"]

    responses = []
    user_ids = [row.user_id for row in users_df.collect()]

    # 将 schema 定义移到循环外部
    schema = StructType([
        StructField("campaign_id", StringType(), True),
        StructField("user_id", StringType(), True),
        StructField("campaign_type", StringType(), True),
        StructField("exposure_time", TimestampType(), True),
        StructField("click_time", TimestampType(), True),
        StructField("conversion_time", TimestampType(), True),
        StructField("conversion_type", StringType(), True),
        StructField("conversion_value", DoubleType(), True)
    ])

    for user_id in user_ids:
        n_campaigns = random.randint(0, 5)
        for i in range(n_campaigns):
            campaign_id = f"campaign_{random.randint(1, 50)}"
            campaign_type = random.choice(campaign_types)

            # 曝光时间
            days_ago = random.randint(0, 30)
            exposure_time = datetime.now() - timedelta(days=days_ago, hours=random.randint(0, 23),
                                                       minutes=random.randint(0, 59))

            # 是否点击
            clicked = random.random() < 0.3
            click_time = exposure_time + timedelta(minutes=random.randint(1, 60)) if clicked else None

            # 是否转化
            converted = clicked and random.random() < 0.5
            conversion_time = click_time + timedelta(minutes=random.randint(1, 360)) if converted else None
            conversion_type = random.choice(conversion_types) if converted else None
            # 关键修改：将默认值从 0 改为 0.0，确保所有 conversion_value 都是浮点数
            conversion_value = round(random.uniform(10, 1000), 2) if converted and conversion_type == "Purchase" else 0.0

            responses.append(Row(
                campaign_id=campaign_id,
                user_id=user_id,
                campaign_type=campaign_type,
                exposure_time=exposure_time,
                click_time=click_time,
                conversion_time=conversion_time,
                conversion_type=conversion_type,
                conversion_value=conversion_value
            ))

    return spark.createDataFrame(responses, schema=schema)

# 生成模拟数据
users = generate_mock_users(1000)
behaviors = generate_mock_behaviors(users, 15)
orders, order_items = generate_mock_orders(users, behaviors)
marketing_responses = generate_mock_marketing_responses(users)

# 保存模拟数据
users.write.csv("../mock_data/user_profile.csv",mode="overwrite")
behaviors.write.csv("../mock_data/user_behavior.csv", mode="overwrite")
orders.write.csv("../mock_data/order_transaction.csv", mode="overwrite")
order_items.write.csv("../mock_data/order_item.csv", mode="overwrite")
marketing_responses.write.csv("../mock_data/marketing_response.csv", mode="overwrite")

print("模拟数据生成完成并保存至mock_data目录")